Skip to content

[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat - #9

Closed
JaredforReal wants to merge 51 commits into
mainfrom
perf/glm53-flash-masked-mha
Closed

JaredforReal wants to merge 51 commits into
mainfrom
perf/glm53-flash-masked-mha

Conversation

@JaredforReal

@JaredforReal JaredforReal commented Sep 6, 2026

Copy link
Copy Markdown
Owner

One of three independent GLM-5.3-Flash perf PRs (#7: decode cleanups; #8: FlashKDA prefill; this one: masked-MHA prefill). Each is benchmarked on its own below. Rebased onto main (previously on wire_sm90_kernels); the one piece it needed from there is carried as its own first commit.

Summary

On main, GLM-5.3-Flash's MLA layout (qk_nope 256, qk_rope 0, v 256) is unknown to the FlashAttention prefill backend and to the masked-MHA allow-list, so MLACommonImpl logs No MLA prefill backend supports this model and every prefill token goes through the per-token top-k MQA kernel. Profiling shows that kernel is KV-gather bound at the HBM roofline: 5.9 ms/layer per 16k-token chunk, and still 1.6 ms/layer for 8x2048 where dense MHA needs ~0.2 ms.

Three commits:

  1. Register the (256, 0, 256) dims with FlashAttnPrefillBackend (7 lines). Same lines as in the SM90 wiring PR [perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA vllm-project/vllm#55385; kept separate so it can be dropped once that lands. This alone gives the dense-MHA prefill path for sequences ≤ index_topk (2048).
  2. Masked-MHA allow-list: _is_masked_mha_available only knew (128, 512, 128, 64, 128) and (64, 512, 192, 64, 256); GLM-5.3-Flash is (64, 512, 256, 0, 256). The kernels are identical for qk_head_dim 256 / v 256, so sequences between 2k and the existing FLASHINFER_MLA_SPARSE thresholds (36k at TP4, 64k at TP8) now use FA4 masked MHA instead of the per-token top-k MQA kernel.
  3. NoPE K concat: _concat_k_nope_k_pe returns k_nope directly when there is no RoPE part instead of allocating and copying a same-sized tensor (134 MB/layer per 16k chunk).

Performance of this PR alone

Measured on top of #7 + #8 (main + #7 + #8) vs the same + this PR, 4x GB300, TP4, --attention-backend FLASHINFER_MLA_SPARSE --max-model-len 69632 --max-num-seqs 256 --max-num-batched-tokens 16384, prefix caching disabled (hit rate 0.0% checked in the server log), vllm bench serve random dataset with warmups, back-to-back in the same session. Decode tok/s is the steady-state window value; TPOT is the per-request median.

point metric +#7 +#8 +#7 +#8 +#9 delta
1k/512 c=1 decode tok/s 147 147 -0.2%
1k/512 c=1 median TPOT ms 6.47 6.49 +0.3%
1k/512 c=64 decode tok/s 3132 3133 +0.0%
1k/512 c=64 median TPOT ms 19.24 19.09 -0.8%
1k/512 c=256 decode tok/s 7717 7755 +0.5%
1k/512 c=256 median TPOT ms 31.30 31.29 -0.0%
prefill 8x2048 mean TTFT ms 395 361 -8.5%
prefill 2x8192 mean TTFT ms 330 315 -4.6%
prefill 8x32768 mean TTFT ms 4693 4347 -7.4%
prefill 2x65536 mean TTFT ms 2450 2365 -3.5%
32k/256 c=16 decode tok/s 316 322 +1.9%
32k/256 c=16 mean TTFT ms 3642 3482 -4.4%
32k/256 c=16 median TPOT ms 37.50 36.97 -1.4%

Prefill up to the masked-MHA threshold (≤ 36k at TP4): −8.5% TTFT for 8x2048 (dense MHA), −7.4% for 8x32k (masked MHA), −4.4% TTFT in the 32k-context serving point. 2x8192 shows −4.6%, but that 16-request point has a ~4% run-to-run noise floor (last column of the ablation below), so treat it as indicative only. 64k prompts are above the threshold and stay on the sparse kernel; the −3.5% there (0.0% noise floor) is the K-concat removal alone. Decode is unchanged (±0.8%).

Ablation of the whole series

Cumulative, same session, same settings: main → +#7 → +#8 → +#9 (this PR) → main again.

point metric main +#7 +#7 +#8 +#7 +#8 +#9 main again
1k/512 c=1 decode tok/s 144 146 (+1.8%) 147 (+2.5%) 147 (+2.3%) 144 (+0.3%)
1k/512 c=1 median TPOT ms 6.57 6.46 (-1.7%) 6.47 (-1.5%) 6.49 (-1.2%) 6.57 (+0.0%)
1k/512 c=64 decode tok/s 2964 3090 (+4.3%) 3132 (+5.7%) 3133 (+5.7%) 2972 (+0.3%)
1k/512 c=64 median TPOT ms 20.02 19.26 (-3.8%) 19.24 (-3.9%) 19.09 (-4.6%) 20.02 (+0.0%)
1k/512 c=256 decode tok/s 7147 7367 (+3.1%) 7717 (+8.0%) 7755 (+8.5%) 7180 (+0.5%)
1k/512 c=256 median TPOT ms 33.87 32.87 (-3.0%) 31.30 (-7.6%) 31.29 (-7.6%) 33.77 (-0.3%)
prefill 8x2048 mean TTFT ms 461 449 (-2.5%) 395 (-14.4%) 361 (-21.7%) 453 (-1.8%)
prefill 2x8192 mean TTFT ms 389 380 (-2.4%) 330 (-15.3%) 315 (-19.2%) 374 (-3.8%)
prefill 8x32768 mean TTFT ms 5258 5108 (-2.8%) 4693 (-10.7%) 4347 (-17.3%) 5258 (+0.0%)
prefill 2x65536 mean TTFT ms 2738 2661 (-2.8%) 2450 (-10.5%) 2365 (-13.6%) 2738 (-0.0%)
32k/256 c=16 decode tok/s 287 296 (+3.1%) 316 (+9.8%) 322 (+11.9%) 288 (+0.2%)
32k/256 c=16 mean TTFT ms 3883 3835 (-1.2%) 3642 (-6.2%) 3482 (-10.3%) 3880 (-0.1%)
32k/256 c=16 median TPOT ms 41.34 40.10 (-3.0%) 37.50 (-9.3%) 36.97 (-10.6%) 41.29 (-0.1%)

Caveat: runtime JIT

The first request in each new sequence-length bucket JIT-compiles the FA4 CuTeDSL mask_mod variant (BlackwellFusedMultiHeadAttentionForward) plus the Triton _scatter_topk_single_req_kernel (~7 s each; jit_monitor warns). FA4MLAPrefillKernel.get_warmup_keys does not cover mask_mod, so this is a pre-existing gap for every masked-MHA model; a warmup for the masked variants should be a follow-up before relying on this in production.

Accuracy (per build, same session)

Each rung of the ablation was also evaluated on its own: gsm8k (1319 questions, 5-shot, greedy, lm_eval local-completions), prompt-logprob agreement on real 4k/12k/30k/60k prompts (mean |Δ logprob| per token vs the main build; the "main again" row is the run-to-run noise floor of this FP8 model), and a needle-in-a-haystack retrieval set (12 codes per length at 6k/16k/30k tokens, greedy, reasoning_effort=low), which exercises the 2k-36k prefill range where #9 switches kernels.

| build | gsm8k flexible-extract % | gsm8k strict-match % | prompt-logprob mean|Δ| vs main (4k/12k/30k/60k) | next token | needle hits (6k/16k/30k) | needle, 3 more passes (concurrency 6) |
|---|---|---|---|---|---|---|
| main | 93.33 ± 0.69 | 93.10 ± 0.70 | (reference) | | 12/12 / 12/12 / 12/12 | 34/36 / 36/36 / 35/36 |
| +#7 | 93.03 ± 0.70 | 92.95 ± 0.71 | 0.040 / 0.055 / 0.051 / 0.108 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 36/36 / 34/36 |
| +#7 +#8 | 93.25 ± 0.69 | 93.18 ± 0.69 | 0.099 / 0.072 / 0.060 / 0.136 | all same | 12/12 / 11/12 / 11/12 | - |
| +#7 +#8 +#9 | 93.10 ± 0.70 | 93.10 ± 0.70 | 0.099 / 0.072 / 0.060 / 0.127 | all same | 12/12 / 12/12 / 11/12 | 35/36 / 34/36 / 36/36 |
| main again | 92.87 ± 0.71 | 92.87 ± 0.71 | 0.041 / 0.055 / 0.051 / 0.118 | all same | 12/12 / 12/12 / 12/12 | - |

Duplicate-work check

gh pr list --repo vllm-project/vllm --state open --search "NoPE masked MHA 256" / "GLM-5.3-Flash prefill": vllm-project#55385 covers the FA dims (commit 1 here is the same 7 lines); vllm-project#54951 / vllm-project#55222 / vllm-project#55543 touch the indexer prefill sharding, the indexer workspace and SM90 fp8 KV. None adds the masked-MHA layout or the K-concat skip.

Tests

  • tests/kernels/test_mhc_kernels.py, tests/kernels/test_glm5next_kda_recurrent_strided.py unaffected; prompt-logprob A/B on real 4k/12k/30k/60k prompts vs main: mean|Δ| 0.06-0.12 vs a main-vs-main noise floor of 0.04-0.11 (≈2.5x at 4k where dense/masked MHA replaces the sparse kernel), identical next tokens; gsm8k 1319 questions (5-shot, greedy): main 93.03% ± 0.70, full series 92.65% ± 0.72 (measured on the three PRs together).
  • pre-commit run ruff-check / ruff-format: passed.
  • AI assistance: developed with Claude Fable 5.1 and reviewed by the submitter; full analysis in ~/notes/glm53flash-perf/REPORT.md, ablation logs in ~/notes/glm53flash-perf/logs/e2e/ablate_*.

peakcrosser7 and others added 10 commits September 5, 2026 14:02
…oject#55375)

Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com>
…3835)

Signed-off-by: WikAlie <194917782+wei-core@users.noreply.github.com>
Co-authored-by: WikAlie <194917782+wei-core@users.noreply.github.com>
Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
Signed-off-by: linitra24 <renshuang.zhou@daocloud.io>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
vllm-project#55448)

Signed-off-by: Luca Motz <luca.motz@icloud.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
…GB PCIe (vllm-project#55511)

Signed-off-by: Baki Burak Ogun <63836730+bakiburakogun@users.noreply.github.com>
…t#51898)

Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com>
Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com>
…hing and spec-decoding (vllm-project#53614)

Signed-off-by: Ziming Huang <zelda.huanghuang@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: levius <2114377220@qq.com>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The changes are narrowly scoped, consistent with existing call patterns, and avoid extra allocation/copy without altering tensor shapes for the NoPE (pe_dim=0) case.

Pull request overview

Enables faster SM100 prefill for GLM-5.3-Flash’s NoPE MLA layout by (1) allowing masked-MHA for the (64, 512, 256, 0, 256) dimension tuple and (2) avoiding an unnecessary K concatenation allocation/copy when the RoPE portion is empty.

Changes:

  • Extend SM100 masked-MHA allow-list to include GLM-5.3-Flash’s NoPE (qk_nope=256, qk_rope=0, v=256) layout.
  • Optimize MLA K construction by returning k_nope directly when k_pe has a zero last-dimension (no RoPE), skipping allocation/copy.
File summaries
File Description
vllm/model_executor/layers/attention/sparse_mla_attention.py Adds GLM-5.3-Flash NoPE dimension tuple to SM100 masked-MHA availability allow-list.
vllm/model_executor/layers/attention/mla_attention.py Skips K concatenation when k_pe is empty to avoid unnecessary allocation/copy in NoPE MLA.
Review details
  • Files reviewed: 2/2 changed files
  • Comments generated: 0
  • Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Alex-ai-future and others added 15 commits September 6, 2026 13:27
…ct#55075)

Signed-off-by: Alex <jihui.huang@daocloud.io>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
…T2 (vllm-project#55461)

Signed-off-by: Wang Zupeng <zupenwang@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
vllm-project#54362)

Signed-off-by: Will Hampson <whamp@ggl.slmail.me>
Co-authored-by: Will Hampson <whamp@ggl.slmail.me>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
vllm-project#54288)

Signed-off-by: almogtavor <almogtavor@gmail.com>
Signed-off-by: Almog Tavor <70065337+almogtavor@users.noreply.github.com>
Co-authored-by: Or Ozeri <or@ozery.com>
…ct#55529)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
…-project#54659)

Signed-off-by: zhouyou9505 <zhouyou9505@gmail.com>
Co-authored-by: Cursor Grok 4.6 <noreply@cursor.com>
…-project#55457)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
…project#55455)

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Codex <noreply@openai.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Signed-off-by: Wenhua Cheng <wenhua.cheng@intel.com>
…uted experts (vllm-project#53161)

Signed-off-by: fai <fangzhouai@gmail.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Cursor Agent <agent@cursor.com>
…llm-project#54975)

Signed-off-by: Big2Wheel <1404207590@qq.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
AndreasKaratzas and others added 15 commits September 7, 2026 06:21
…ect#55642)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
…rEngine (vllm-project#41567)

Signed-off-by: Teng Ma <sima.mt@alibaba-inc.com>
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Signed-off-by: Zhou ziheng <jiaranran2@gmail.com>
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Tianyu Guo <guoty@inferact.ai>
Co-authored-by: Zhou ziheng <jiaranran2@gmail.com>
Co-authored-by: jiangkuaixue123 <jiangxiaozhou111@163.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
…oject#55653)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
…llm-project#54643)

Signed-off-by: zhewenl <zhewenl@users.noreply.github.com>
Co-authored-by: zhewenl <zhewenl@users.noreply.github.com>
…r BF16 variant (vllm-project#53586)

Signed-off-by: Zihua Wu <zihuaw@nvidia.com>
Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>
…ternal_tokens (vllm-project#47505)

Signed-off-by: Alex <alex.tech.lab@outlook.com>
Signed-off-by: Alex-ai-future <Alex-ai-future@users.noreply.github.com>
…latforms (vllm-project#55660)

Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
…tention prefill backend

GLM-5.3-Flash (qk_nope 256, qk_rope 0, v 256) was not in
FlashAttnPrefillBackend.supports_mla_dimensions, so sparse MLA fell back to
"No MLA prefill backend supports this model" and every prefill token went
through the per-token top-k MQA kernel. The kernels for qk_head_dim 256 /
v_head_dim 256 are the ones already used for the (192, 64, 256) layout.

Same 7 lines as in the SM90 wiring PR (vllm-project#55385); kept as a
separate commit so it can be dropped once that lands.

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Jared Wen <jaredwen@inferact.ai>
…4, 512, 256, 0, 256) layout

wire_sm90_kernels registers the (256, 0, 256) MLA dimensions with the
FlashAttention prefill backend, which gives GLM-5.3-Flash the dense-MHA
prefill path for sequences up to index_topk (2048). Sequences above that
still went through the per-token top-k MQA kernel on SM100, which is
KV-gather bound (5.9 ms/layer per 16k-token chunk, ~4.8 ms of it pure
HBM traffic), because the masked-MHA allow-list only knew the
(128, 512, 128, 64, 128) and (64, 512, 192, 64, 256) layouts. The masked
kernel path is identical for qk_head_dim 256 / v 256, so add the GLM
layout; the existing FLASHINFER_MLA_SPARSE TP4/TP8 thresholds (masked MHA
up to 36k / 64k tokens) apply.

Prefill TTFT on 4x GB300 (TP4, together with the FlashKDA and decode
cleanup PRs): 8x2048 462 -> 363 ms, 2x8192 382 -> 317 ms, 32k/256 c16
TTFT 3876 -> 3425 ms. First use of a new sequence-length bucket JIT
compiles the FA4 mask_mod variant (~7 s); the FA4 MLA prefill warmup does
not cover mask_mod yet.

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Jared Wen <jaredwen@inferact.ai>
With qk_rope_head_dim == 0 the concatenated K is just k_nope; return it
instead of allocating and copying a same-sized tensor (134 MB per layer
for a 16k-token chunk on GLM-5.3-Flash).

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Jared Wen <jaredwen@inferact.ai>
…project#54404)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
…rror (vllm-project#50257)

Signed-off-by: shaolila <shaolila@buaa.edu.cn>
Co-authored-by: shaolila <shaolila@buaa.edu.cn>
…ect#53580)

Signed-off-by: Marceli Fylcek <marceli.fylcek@intel.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Chaojun Zhang <chaojun.zhang@intel.com>
Signed-off-by: Sean Westfall <sean.westfall@gmail.com>
@JaredforReal JaredforReal changed the title [Perf][GLM-5.3-Flash] Masked-MHA sparse prefill for the NoPE (256, 0, 256) layout on SM100 + skip the NoPE K concat [Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat Sep 7, 2026
@JaredforReal
JaredforReal force-pushed the perf/glm53-flash-masked-mha branch from 5299590 to 85928b9 Compare September 7, 2026 10:39
@JaredforReal
JaredforReal changed the base branch from wire_sm90_kernels to main September 7, 2026 10:39
gau-nernst and others added 4 commits September 7, 2026 12:15
…54890)

Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
…LE spec decode (vllm-project#52771)

Signed-off-by: Kam Basra <kameldipbasra@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
…50195)

Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
@JaredforReal

Copy link
Copy Markdown
Owner Author

Upstream draft: vllm-project#55738

mergify Bot and others added 2 commits September 7, 2026 13:24
Signed-off-by: Jared Wen <w13431838023@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.